use crate::types::TResult;
use image::{DynamicImage, GenericImage, GenericImageView, ImageBuffer, Pixel, Primitive};
use xcap::Monitor;

/// 目标屏幕
#[derive(Clone, Copy)]
pub enum ScreenTarget {
    /// 主屏幕
    Primary,
    /// 所有屏幕
    All,
}

/// 获取主屏幕
pub fn get_primary() -> TResult<ImageBuffer<image::Rgba<u8>, Vec<u8>>> {
    Ok(Monitor::from_point(0, 0)?.capture_image()?)
}

/// 获取所有屏幕
pub fn get_all() -> TResult<Vec<ImageBuffer<image::Rgba<u8>, Vec<u8>>>> {
    let monitors = Monitor::all()?;
    let mut imgs = vec![];
    for monitor in monitors {
        imgs.push(monitor.capture_image()?);
    }
    Ok(imgs)
}

/// 横向拼接
fn h_concat<I, P, S>(imgs: &[I]) -> TResult<ImageBuffer<P, Vec<S>>>
where
    I: GenericImageView<Pixel = P>,
    P: Pixel<Subpixel = S> + 'static,
    S: Primitive + 'static,
{
    // The final width is the sum of all images width.
    let img_width_out: u32 = imgs.iter().map(|img| img.width()).sum();

    // The final height is the maximum height from the input images.
    let img_height_out: u32 = imgs.iter().map(|img| img.height()).max().unwrap_or(0);

    // Initialize an image buffer with the appropriate size.
    let mut img_buf = image::ImageBuffer::new(img_width_out, img_height_out);
    let mut accumulated_width = 0;

    // Copy each input image at the correct location in the output image.
    for img in imgs {
        img_buf.copy_from(img, accumulated_width, 0)?;
        accumulated_width += img.width();
    }

    Ok(img_buf)
}

/// 获取所有屏幕截图并拼接成一张
pub fn get_all_to_one() -> TResult<ImageBuffer<image::Rgba<u8>, Vec<u8>>> {
    let imgs = get_all()?;
    h_concat(&imgs)
}

/// 截取屏幕
pub fn screenshot_to_img(screen: ScreenTarget) -> TResult<DynamicImage> {
    let img = match screen {
        ScreenTarget::Primary => get_primary()?,
        ScreenTarget::All => get_all_to_one()?,
    };
    Ok(DynamicImage::ImageRgba8(img))
}

/// 获取屏幕区域截图
pub fn capture_region(x: i32, y: i32, width: u32, height: u32) -> TResult<DynamicImage> {
    let img = get_primary()?;
    let cropped = img.view(x as u32, y as u32, width, height).to_image();
    Ok(DynamicImage::ImageRgba8(cropped))
}